package com.example.springboot_sdse_air.service;

import com.example.springboot_sdse_air.common.OneApiConfig;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.netty.http.client.HttpClient;

import javax.annotation.Resource;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeoutException;

/**
 * 功能：
 * 作者：王彬 202200111149
 * 日期：2025-04-29 9:23
 */
@Service
public class DeepSeekService {

    private final WebClient webClient;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final OneApiConfig oneApiConfig;
    private final String apiBase;
    private final String apiKey;
    private final String model;

    @Autowired
    public DeepSeekService(OneApiConfig oneApiConfig) {
        this.oneApiConfig = oneApiConfig;
        apiBase = this.oneApiConfig.getUrl();
        apiKey = this.oneApiConfig.getKey();
        model = this.oneApiConfig.getModel();

        this.webClient = WebClient.builder()
                .baseUrl(apiBase)
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .build();
    }

    public Flux<String> streamResponse(String prompt) {
        Map<String, Object> requestBody = createRequestBody(prompt);

        return webClient.post()
                .uri("/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.TEXT_EVENT_STREAM)
                .bodyValue(requestBody)
                .retrieve()
                .bodyToFlux(String.class)
                .filter(responseJson -> !"[DONE]".equals(responseJson.trim()))  // 过滤掉终止标记
                .map(responseJson -> {
                    try {
                        // 解析响应 JSON
                        JsonNode root = objectMapper.readTree(responseJson);
                        JsonNode choices = root.path("choices");
                        if (choices.isArray() && choices.size() > 0) {
                            // 注意使用 delta.content
                            System.out.println(choices.get(0).path("delta").path("content").asText());
                            return choices.get(0).path("delta").path("content").asText();
                        } else {
                            return "[无法解析响应内容]";
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                        return "[解析出错]";
                    }
                });
    }

    public Mono<String> fullResponse(String prompt) {
        long start = System.currentTimeMillis();
        Map<String, Object> requestBody = createRequestBody(prompt);

        return webClient.post()
                .uri("/chat/completions")
                .contentType(MediaType.APPLICATION_JSON)
                .accept(MediaType.APPLICATION_JSON)
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(String.class)  // 获取完整响应
                .doOnNext(
                        result -> {
                            System.out.println("收到响应内容：" + result);
                            long end = System.currentTimeMillis();
                            System.out.println("请求 DeepSeek API 总耗时：" + (end - start) + " ms");
                            System.out.println("收到响应内容：" + result);
                        }
                )
                .flatMap(responseJson -> {
                    try {
                        // 解析响应 JSON
                        JsonNode root = objectMapper.readTree(responseJson);
                        JsonNode choices = root.path("choices");
                        if (choices.isArray() && choices.size() > 0) {
                            // 获取 content 字段
                            return Mono.just(choices.get(0).path("message").path("content").asText());
                        } else {
                            return Mono.just("没有有效内容");  // 如果没有解析到有效的 content
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                        return Mono.just("解析错误");  // 如果解析错误
                    }
                })
                .timeout(Duration.ofSeconds(120))  // 控制 WebClient 自身的最大处理时间
                .onErrorResume(TimeoutException.class, e -> Mono.just("请求超时，请稍后重试"));
    }


    private Map<String, Object> createRequestBody(String prompt) {
        Map<String, Object> message_body = new HashMap<>();
        message_body.put("role","user");
        message_body.put("content",prompt);

        List<Map<String, Object>> message = new ArrayList<>();
        message.add(message_body);

        Map<String, Object> body = new HashMap<>();
        body.put("model", model);
        body.put("messages", message);
        body.put("temperature",0.7);
        body.put("max_tokens", 1024);
        body.put("stream", true);

        return body;
    }


//    public Mono<String> testResponse(String prompt) {
//        Map<String, Object> requestBody = createRequestBody(prompt);
//
//        return webClient.post()
//                .uri("/chat/completions")
//                .bodyValue(requestBody)
//                .retrieve()
//                .bodyToMono(String.class)  // 一次性拿整个响应
//                .doOnNext(result -> System.out.println("收到响应内容：" + result))
//                .map(responseJson -> {
//                    try {
//                        JsonNode root = objectMapper.readTree(responseJson);
//                        JsonNode choices = root.path("choices");
//                        if (choices.isArray() && choices.size() > 0) {
//                            System.out.println(choices.get(0).path("message").path("content").asText());
//                            return choices.get(0).path("message").path("content").asText();
//                        } else {
//                            return "[无法解析响应内容]";
//                        }
//                    } catch (Exception e) {
//                        e.printStackTrace();
//                        return "[解析出错]";
//                    }
//                });
//    }
}
